#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import fire

from pymongo import MongoClient


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def GetPossibility(ref: str, alt: str, ances: str, geno: dict):
    total_individual_number = 0
    total_alt_number = 0

    allele = ""
    if ances == alt:
        allele = ref
    else:
        allele = alt

    for item in geno.keys():
        total_individual_number += geno[item]
        count_alt = item.count(allele)
        total_alt_number += count_alt * geno[item]

    total_individual_number = total_individual_number * 2

    return ([total_alt_number, total_individual_number])


def UpdatePopRegionAAF():
    client_genotypes, collection_genotypes = InitMongoDB("localhost", 27017, "hanvcf", "genotypes")

    count = 0

    for result in collection_genotypes.find():
        if "geno" in result:
            pop_aaf = GetPossibility(result["ref"], result["alt"], result["ances"], result["geno"])
            collection_genotypes.update_one({"_id": result["_id"]}, {
                "$set": {"pop_aaf": {"alt_number": pop_aaf[0], "individual_number": pop_aaf[1]}}})
        else:
            collection_genotypes.update_one({"_id": result["_id"]}, {
                "$set": {"pop_aaf": {"alt_number": 0, "individual_number": 0}}})
        if "region_aaf" not in result:
            region_aaf = [0, 0]
            for region_result in collection_genotypes.find(
                    {"chrom": result["chrom"], "pos": result["pos"], "ref": result["ref"], "alt": result["alt"],
                     "genetic_region": result["genetic_region"]}):
                if "geno" in region_result:
                    cpop_aaf = GetPossibility(region_result["ref"], region_result["alt"], region_result["ances"],
                                              region_result["geno"])
                    region_aaf[0] += cpop_aaf[0]
                    region_aaf[1] += cpop_aaf[1]
            collection_genotypes.update_many(
                {"chrom": result["chrom"], "pos": result["pos"], "ref": result["ref"], "alt": result["alt"],
                 "genetic_region": result["genetic_region"]},
                {"$set": {"region_aaf": {"alt_number": region_aaf[0], "individual_number": region_aaf[1]}}})
        count += 1
        if count % 100000 == 0:
            print(str(count) + " records processed.")
    CloseMongoDB(client_genotypes)


def UpdateSNVAAF():
    client_snvs, collection_snvs = InitMongoDB("localhost", 27017, "hanvcf", "snvs")
    client_genotypes, collection_genotypes = InitMongoDB("localhost", 27017, "hanvcf", "genotypes")

    count = 0

    for result in collection_snvs.find():
        snv_aaf = [0, 0]
        for pop_result in collection_genotypes.find(
                {"chrom": result["chrom"], "pos": result["pos"], "ref": result["ref"], "alt": result["alt"]}):
            if "geno" in pop_result:
                cpop_aaf = GetPossibility(pop_result["ref"], pop_result["alt"], pop_result["ances"], pop_result["geno"])
                snv_aaf[0] += cpop_aaf[0]
                snv_aaf[1] += cpop_aaf[1]
        collection_snvs.update_one({"_id": result["_id"]}, {
            "$set": {"snv_aaf": {"alt_number": snv_aaf[0], "individual_number": snv_aaf[1]}}})
        count += 1
        if count % 10000 == 0:
            print(str(count) + " records processed.")
    CloseMongoDB(client_snvs)
    CloseMongoDB(client_genotypes)


def UpdateRegion():
    htfile = open("province_region.txt", "r")
    province_region = {}
    for line in htfile:
        line = line.rstrip()
        splitline = line.split("\t")
        province_region[splitline[0]] = splitline[1]
    htfile.close()

    client_genotypes, collection_genotypes = InitMongoDB("localhost", 27017, "hanvcf", "genotypes")

    for result in collection_genotypes.find():
        collection_genotypes.update_one({"_id": result["_id"]},
                                        {"$set": {"genetic_region": province_region[result["pop_name"]]}})

    CloseMongoDB(client_genotypes)


if __name__ == '__main__':
    # UpdatePopPossibilities()
    # UpdateRegion()
    # UpdatePopRegionAAF()
    fire.Fire()
